import os
import sys
import torch
import pickle
import shutil


def create_directories(base_dir, sub_dirs):
    for sub_dir in sub_dirs:
        os.makedirs(os.path.join(base_dir, sub_dir), exist_ok=True)


def save_image(image_path, dir, original=False):
    file_name = os.path.basename(image_path)
    dir_prefix = '/'.join(dir.split('/')[:5])
    # orginal控制是否取原图
    if not original:
        image_path = dir_prefix + '/inference/images/' + file_name

    os.makedirs(dir, exist_ok=True)
    target_path = os.path.join(dir, file_name)

    shutil.copy(image_path, str(target_path))
    # print(f"{file_name}图片已保存到 {target_path}")


def judge_pos_neg(tensor):
    return any((tensor == 1).bool())


def process_detection(detDataSample, base_result_dir, res_dict, unSure_score_thr):
    """
    detDataSample: dict
     {'img_path': '/home/jacy/GDW/data_0508/labelme_to_coco/val2017/OK (51).jpg',
     'pad_param': array([ 0.,  0., 26., 26.], dtype=float32),
     'ori_shape': (111, 102),
     'batch_input_shape': (640, 640),
     'scale_factor': (5.764705882352941, 5.7657657657657655),
     'img_id': 0,
     'pad_shape': (640, 640),
     'img_shape': (640, 640, 3),
     'ignored_instances': {'labels': tensor([], dtype=torch.int64), 'bboxes': tensor([], size=(0, 4))},
     'gt_instances': {'labels': tensor([0]), 'bboxes': tensor([[23., 17., 93., 90.]])}, # 0为neg，1为pos
     'pred_instances': {'labels': tensor([0, 1, 0, 0, 1]),
                        'scores': tensor([0.7536, 0.1378, 0.0025, 0.0013, 0.0011]),
                        'bboxes': tensor([[ 23.4096,  18.9176,  89.7938,  85.9218],
                                            [ 23.8616,  19.3099,  89.2531,  86.0639],
                                            [  0.0000,  18.2759,  67.0084, 107.9124],
                                            [ 18.6911, 100.8391,  90.0358, 111.0000],
                                            [  0.0000,  16.6500,  63.1221, 110.0392]])}

     }
    """
    # print(detDataSample['pred_instances']['scores'])
    pred_scores = detDataSample['pred_instances']['scores']  # 为torch.tensor类型
    has_greater_than_thr = (pred_scores >= unSure_score_thr).any()
    # print(detDataSample['img_path'])
    if has_greater_than_thr:
        # 处理有推理框的相关照片
        # 取出gt标签
        gt_label_tensor = detDataSample['gt_instances']['labels']
        gt_label = judge_pos_neg(gt_label_tensor)
        # 取出推理的标签
        inf_label_tensor = detDataSample['pred_instances']['labels'][pred_scores >= unSure_score_thr]
        inf_label = judge_pos_neg(inf_label_tensor)
        # 开始分情况copy图片
        if gt_label and inf_label:
            img_path = detDataSample['img_path']
            save_image(img_path, os.path.join(base_result_dir, 'TP'))
            res_dict['TP'] += 1
        elif (not gt_label) and (not inf_label):
            img_path = detDataSample['img_path']
            save_image(img_path, os.path.join(base_result_dir, 'TN'))
            res_dict['TN'] += 1
        elif gt_label and (not inf_label):
            # fp情况
            img_path = detDataSample['img_path']
            save_image(img_path, os.path.join(base_result_dir, 'FP'))
            res_dict['FP'] += 1
        elif (not gt_label) and inf_label:
            # fn情况
            img_path = detDataSample['img_path']
            save_image(img_path, os.path.join(base_result_dir, 'FN'))
            res_dict['FN'] += 1
    else:
        # 将没有推理框的图片放于FN文件夹
        img_path = detDataSample['img_path']
        save_image(img_path, os.path.join(base_result_dir, 'UNSURE'))


def save_accuracy_to_md(result_dict, dir):
    accuracy = (result_dict['TP'] + result_dict['TN']) / sum(result_dict.values())
    recall = result_dict['TP'] / (result_dict['TP'] + result_dict['FN'])
    precision = result_dict['TP'] / (result_dict['TP'] + result_dict['FP'])

    file_name = f"accuracy_{accuracy:.5f}.md"
    os.makedirs(dir, exist_ok=True)
    file_path = os.path.join(dir, file_name)

    # 写入 Markdown 文件
    with open(file_path, 'w') as f:
        f.write(f"# Accuracy: {accuracy:.5f}\n")
        f.write(f"# Recall: {recall:.5f}\n")
        f.write(f"# Precision: {precision:.5f}\n")
        f.write(f"# Result dict: {result_dict}")


def main(args):
    base_result_dir = f"/home/jacy/GDW/data_{args[1]}/result"

    sub_dirs = ['TP', 'TN', 'FP', 'FN', 'UNSURE']
    create_directories(base_result_dir, sub_dirs)

    with open(f'/home/jacy/GDW/data_{args[1]}/inference/pending_result.pkl', 'rb') as f:
        data = pickle.load(f)

    # print(len(data))

    # 记录指标
    res_dict = {'TP': 0, 'TN': 0, 'FP': 0, 'FN': 0}
    # 和test.py的trigger_visualization_hook的阙值保持一致， unSure阙值0.3
    # 在trigger_visualization_hook中低于此值，不会在图片中给出推理框
    unSure_score_thr = 0.3
    for detDataSample in data:
        process_detection(detDataSample, base_result_dir, res_dict, unSure_score_thr)

    save_accuracy_to_md(res_dict, base_result_dir)
    print(f"FP, FN, TP, TN图片以及accuracy已经保存至{base_result_dir}")


if __name__ == "__main__":
    main(sys.argv)
